library("foreign")
library("magrittr")
library("dplyr")
library("broom")
library("stringr")
library("knitr")

balance.function <- function(experiment, subset = "!is.na(data$treat_ind)"){
  data <- read.dta(experiment) #Read in the .dta file.
  t0.covariates.name <- unlist(strsplit(as.character(subset(master.sheet, master.sheet$Data == experiment, 
                                                            xvars)), " ")) #Read in the covariates for balance.
  t0.covariates.name <- gsub('c\\(\\"', "", t0.covariates.name)
  t0.covariates.name <- gsub('\n', "", t0.covariates.name)
  t0.covariates.name <- gsub('\"', "", t0.covariates.name)
  t0.covariates.name <- gsub(',', "", t0.covariates.name)
  t0.covariates.name <- gsub('\\)', "", t0.covariates.name)
  t0.covariates.name <- unique(t0.covariates.name) #De-dup variable names.
  x <- data[,c(t0.covariates.name)] #Make a nice matrix.
  x <- as.matrix(x)
  fixed_digits <- function(xs, n = 2) {
    formatC(xs, digits = n, format = "f")
  }
  # Don't print leading zero on bounded numbers.
  remove_leading_zero <- function(xs) {
    # Problem if any value is greater than 1.0
    digit_matters <- xs %>% as.numeric %>%
      abs %>% is_greater_than(1)
    if (any(digit_matters)) {
      warning("Non-zero leading digit")
    }
    str_replace(xs, "^(-?)0", "\\1")
  }
  # Print three digits of a p-value, but use
  # the "< .001" notation on tiny values.
  format_pval <- function(ps, html = FALSE) {
    tiny <- ifelse(html, "&lt;&nbsp;.001", "< .001")
    ps_chr <- ps %>% fixed_digits(3) %>%
      remove_leading_zero
    ps_chr[ps < 0.001] <- tiny
    ps_chr
  }
 # m_table$p.value %>% format_pval
  fix_names <- . %>%
    str_replace(".Intercept.", "Intercept") %>%
    str_replace("text = subset", "") %>%
    str_replace("x\\[", "") %>%
    str_replace("which\\(eval\\(parse\\(\\)\\)\\), ]", "")
  two_digits <- . %>% fixed_digits(2)
  table_names <- c("Parameter", "Estimate", "SE",
                   "_t_", "_p_")
  lm.table <- lm(data[which(eval(parse(text = subset))),]$treat_ind ~ x[which(eval(parse(text = subset))),])
  f <- summary(lm.table)$fstatistic
  p.fstat <- round(pf(f[1], f[2], f[3], lower=FALSE), digits = 3)
  
  format_model_table <- . %>%
    mutate_each(funs(two_digits),
                -term, -p.value) %>%
    mutate(term = fix_names(term),
           p.value = format_pval(p.value)) %>%
    set_colnames(table_names)
  lm.table %>%
    tidy %>%
    format_model_table %>%
    kable(align = "r", caption = paste0("Test of covariate balance. F-statistic from this multivariate regression is ", p.fstat, "."))
}
